#pragma once

#include "Rdma.h"
#include "MemoryManager.h"

#include <stdlib.h>

#include <unordered_map>
#include <iostream>
#include <deque>

struct RouteInf
{
  uint32_t self_id;
  uint32_t r_key;
  uint32_t qpn;
  uint16_t lid;
  uint8_t  gid[16];
  uint64_t ring_begin;
  uint64_t ring_end;
}__attribute__((packed));

enum SessErrors
{
    OK = 0,
    WRITE_ROUTE_FAILED = -1,
    READ_ROUTE_FAILED = -2,
    MODIFY_QP_FAILED = -3,
    CREATE_QP_FAILED = -4,
    DESTROY_QP_FAILED = -5,
    ALLOC_RING_BUF_FAILED = -6,
    FREE_RING_BUF_FAILED = -7,
    TCP_ESTABLISHED_FAILED = -8,
    OTHER_ERROR = -9,
};

typedef void (*RpcSessHandler)(void* rpc, void* user_ctx);

enum SessStatus
{
    CONNECTED = 0,
    WAITING_EST, // wait to establish connection
    WAITING_CLOSE, // wait to close connection
};
class Session {
private:
    struct RingBuf{
        uint64_t begin;
        uint64_t end;
    };
    struct EpollEventCtx {
        int fd;
        Session* sess;
    };
    MemoryManager* mm_;
    Rdma* rdma_res_;
    struct ibv_qp* qp_;
    int sock_;

    uint32_t peer_rkey_;
    struct RingBuf self_rbuf_;
    struct RingBuf peer_rbuf_;

    uint64_t head_copy_;
    uint64_t tail_copy_;
    uint32_t processed_;
    bool is_full_;

    uint32_t local_recv_num_ = 0;

    uint64_t self_tail_;
    uint64_t self_head_;

    uint32_t self_id_;
    uint32_t peer_id_;

    uint32_t inflight_;
    uint32_t credits_;
    uint16_t status_;
    RpcSessHandler close_cb_;
    void* close_cb_arg_;




    std::deque<void*> batch_queue_;


public:
    Session(MemoryManager* mm, struct Rdma* res);
    ~Session();
    int Init(uint32_t& cur_sess_id);
    int Destroy();
    int SessConnect(const char* remote_uri);
    int SessListen(const char* local_uri, int& efd);

    int SessAccept(int efd, std::unordered_map<uint32_t, Session*>& id2sess, MemoryManager* mm, Rdma* res, uint32_t& cur_sess);
    
    int SessClose();

    int SessReadRoute(int fd);
    int SessWriteRoute(int fd);
    

    void* deque_sess()
    {
        void *ret = batch_queue_.front();
        batch_queue_.pop_front();
        return ret;
    }
    void enque_req(void* ex)
    {
        batch_queue_.push_back(ex);
    }
    void* queue_front() {return batch_queue_.front();}
    decltype(batch_queue_.size()) queue_size() const {return batch_queue_.size();}
    inline struct ibv_qp* get_qp() const { return qp_; }
    inline void set_self_id(uint32_t id) { self_id_ = id; }
    inline uint32_t get_self_id() const { return self_id_; }
    inline void set_peer_id(uint32_t id) { peer_id_ = id; }
    inline uint32_t get_peer_id() const { return peer_id_; }
    inline void set_close_cb(RpcSessHandler cb, void* arg) 
    { 
        close_cb_ = cb;
        close_cb_arg_ = arg;  
    }

    inline void set_sess_status(uint16_t stat) { status_ = stat; }
    inline uint16_t get_sess_status() const { return status_; }
    inline void add_inflight_ops(uint32_t adds) { inflight_+=adds; }
    inline void sub_inflight_ops(uint32_t subs) { inflight_-=subs; }

    inline uint32_t get_inflight_ops() const { return inflight_; }
    
    uint64_t get_peer_buf(uint64_t len);
    void free_peer_buf(uint32_t len);
    uint64_t poll_msg();
    uint32_t free_msg(uint64_t msg);

    inline uint32_t get_processed() const {return processed_;}
    inline void add_processed(uint32_t num) { processed_ += num; } 
    inline void sub_processed(uint32_t num) { processed_ -= num; } 
    inline uint32_t get_peer_rkey() const {return peer_rkey_;}
    inline uint64_t get_peer_beg() const {return peer_rbuf_.begin;}
    inline uint64_t get_peer_end() const {return peer_rbuf_.end;}
};